from scipy.optimize import linear_sum_assignment
import numpy as np
import tensorflow as tf
from tensorflow.keras.datasets import cifar10
from DML import triplet_loss

# # 启用内存增长
# gpus = tf.config.experimental.list_physical_devices('GPU')
# if gpus:
#     try:
#         for gpu in gpus:
#             tf.config.experimental.set_memory_growth(gpu, True)
#     except RuntimeError as e:
#         print(e)


# 加载CIFAR-10数据集作为示例
(X_train, y_train), (X_test, y_test) = cifar10.load_data()
X_test = X_test.astype('float32') / 255
y_test = y_test.flatten()

# 加载保存的模型
model = tf.keras.models.load_model('src/model_pth/triplet_model.keras', custom_objects={'triplet_loss': triplet_loss})

def optimal_transport_distance(embeddings, epsilon=0.01):
    n = embeddings.shape[0]
    C = np.zeros((n, n))
    
    for i in range(n):
        for j in range(n):
            C[i, j] = np.linalg.norm(embeddings[i] - embeddings[j])
    
    P = np.exp(-C / epsilon)
    row_sums = P.sum(axis=1)
    P /= row_sums[:, np.newaxis]
    
    return P

# 获取嵌入向量
# 将测试数据分成三个部分
X_test_anchor = X_test
X_test_positive = X_test
X_test_negative = X_test

# 预测嵌入向量
embeddings = model.predict([X_test_anchor, X_test_positive, X_test_negative])

# 提取嵌入向量
anchor_embeddings = embeddings[0]
positive_embeddings = embeddings[1]
negative_embeddings = embeddings[2]

# 计算最优传输距离
P = optimal_transport_distance(anchor_embeddings)
# print(P)

# 模型评估
def evaluate_model(embeddings, y_test):
    n = embeddings.shape[0]
    correct = 0
    
    for i in range(n):
        for j in range(i+1, n):
            if y_test[i] == y_test[j]:
                if np.linalg.norm(embeddings[i] - embeddings[j]) < 1.0:
                    correct += 1
            else:
                if np.linalg.norm(embeddings[i] - embeddings[j]) > 1.0:
                    correct += 1
    
    accuracy = correct / (n * (n - 1) / 2)
    return accuracy

accuracy = evaluate_model(anchor_embeddings, y_test)
print(f"Model Accuracy: {accuracy:.4f}")